##########################################################################################

library(data.table)
library(optparse)
library(dplyr)
library(parallel)
library(ggplot2)
library(ggpubr)
library(cowplot)

##########################################################################################

option_list <- list(
    make_option(c("--tcga_rate_file"), type = "character") ,
    make_option(c("--njmu_rate_file"), type = "character") ,
    make_option(c("--tumor_type"), type = "character") ,
    make_option(c("--mutsig_list_file"), type = "character") ,
    make_option(c("--early_list_file"), type = "character") ,
    make_option(c("--show_file"), type = "character") ,
    make_option(c("--out_path"), type = "character")
)

if(1!=1){

    tumor_type <- "STAD"
    tcga_rate_file <- "~/20220915_gastric_multiple/dna_combine/public_ref/TCGA/"
    njmu_rate_file <- "/public/home/xxf2019/20220205_lungSomatic/Analysis_LUAD/maf/LUAD.mut_rate.tsv"
    out_path <- "/public/home/xxf2019/20220205_lungSomatic/Analysis_LUAD/images/Figures/TCGA_Compare/BaseLineSmoke_white"
    show_file <- "/public/home/xxf2019/20220205_lungSomatic/Analysis_LUAD/mutsig_check/show_gene.list"

    mutsig_list_file <- "/public/home/xxf2019/20220205_lungSomatic/Analysis_LUAD/mutsig_check/mutsig_gene_check.txt"
    early_list_file <- "/public/home/xxf2019/20220205_lungSomatic/Analysis_LUAD/mutsig_check/Early_Enrichment.txt"

}

###########################################################################################

parseobj <- OptionParser(option_list=option_list, usage = "usage: Rscript %prog [options]")
opt <- parse_args(parseobj)
print(opt)

tcga_rate_file <- opt$tcga_rate_file
njmu_rate_file <- opt$njmu_rate_file
mutsig_list_file <- opt$mutsig_list_file
early_list_file <- opt$early_list_file
tumor_type <- opt$tumor_type
show_file <- opt$show_file
out_path <- opt$out_path

###########################################################################################

dir.create(out_path , recursive = T)

###########################################################################################

Variant_Type <- c("Missense_Mutation","Nonsense_Mutation","Frame_Shift_Ins","Frame_Shift_Del","In_Frame_Ins","In_Frame_Del","Splice_Site","Nonstop_Mutation")
time_order <- c("Clonal (early)" ,  "Clonal (unsp.)" , "Clonal (late)" ,"Subclonal")
col <- c( "#4DAF4A" , "#377EB8" , "#984EA3" , "#E41A1C")
names(col) <- time_order

###########################################################################################

tcga_mut_rate <- fread(tcga_rate_file,sep = "\t",header = T)
njmu_mut_rate <- fread(njmu_rate_file,sep = "\t",header = T)
show_gene <- data.frame(fread(show_file , header = F))$V1
mutsig_list <- data.frame(fread(mutsig_list_file , header = F))
early_list <- data.frame(fread(early_list_file , header = F))

###########################################################################################
## 提取SMG
njmu_mut_rate <- subset( njmu_mut_rate , gene %in% show_gene )
tcga_mut_rate <- subset( tcga_mut_rate , gene %in% show_gene )

colnames(njmu_mut_rate)[2:ncol(njmu_mut_rate)] <- paste0( colnames(njmu_mut_rate)[2:ncol(njmu_mut_rate)] , "_njmu" )
colnames(tcga_mut_rate)[2:ncol(tcga_mut_rate)] <- paste0( colnames(tcga_mut_rate)[2:ncol(tcga_mut_rate)] , "_tcga" )

## 比较突变率
dat_rate <- merge(njmu_mut_rate , tcga_mut_rate , by="gene" , all.y=T)


###########################################################################################
mutsig_gene <- mutsig_list$V1
early_gene <- early_list$V1


index <- which(dat_rate$gene %in% mutsig_gene & dat_rate$gene %in% early_gene)
dat_rate$gene[index] <- paste0( "*" , dat_rate$gene[index] , "*" )
index <- which( dat_rate$gene %in% early_gene )
dat_rate$gene[index] <- paste0( dat_rate$gene[index] , "*" )
index <- which( dat_rate$gene %in% mutsig_gene )
dat_rate$gene[index] <- paste0( "*" , dat_rate$gene[index])


###########################################################################################
## 计算显著性
result1 <- Reduce(function(x,y)bind_rows(x,y),mclapply(dat_rate$gene,function(geneN){
    print(geneN)

    tmp <- subset(dat_rate , gene == geneN)

    ## 总体突变情况
    a <- tmp$freq_njmu
    c <- round(tmp$freq_njmu/tmp$rate_njmu) - tmp$freq_njmu
    b <- tmp$freq_tcga
    d <- round(tmp$freq_tcga/tmp$rate_tcga) - tmp$freq_tcga

    if(is.na(a)){a=0}
    if(is.na(b)){b=0}
    if(is.na(c)){c=0}
    if(is.na(d)){d=0}

    result=fisher.test(matrix(c(a,b,c,d),nrow=2))

    p=result[["p.value"]]
    OR=round(result[["estimate"]][["odds ratio"]],3)

    tmp$p_All <- p
    tmp$OR_All <- OR

    ## 吸烟
    a <- tmp$freq_smoke_njmu
    c <- round(tmp$freq_smoke_njmu/tmp$rate_smoke_njmu) - tmp$freq_smoke_njmu
    b <- tmp$freq_smoke_tcga
    d <- round(tmp$freq_smoke_tcga/tmp$rate_smoke_tcga) - tmp$freq_smoke_tcga

    if(is.na(a)){a=0}
    if(is.na(b)){b=0}
    if(is.na(c)){c=0}
    if(is.na(d)){d=0}

    result=fisher.test(matrix(c(a,b,c,d),nrow=2))

    p=result[["p.value"]]
    OR=round(result[["estimate"]][["odds ratio"]],3)

    tmp$p_Smoke <- p
    tmp$OR_Smoke <- OR

    ## 非吸烟
    a <- tmp$freq_nosmoke_njmu
    c <- round(tmp$freq_nosmoke_njmu/tmp$rate_nosmoke_njmu) - tmp$freq_nosmoke_njmu
    b <- tmp$freq_nosmoke_tcga
    d <- round(tmp$freq_nosmoke_tcga/tmp$rate_nosmoke_tcga) - tmp$freq_nosmoke_tcga

    if(is.na(a)){a=0}
    if(is.na(b)){b=0}
    if(is.na(c)){c=0}
    if(is.na(d)){d=0}

    result=fisher.test(matrix(c(a,b,c,d),nrow=2))

    p=result[["p.value"]]
    OR=round(result[["estimate"]][["odds ratio"]],3)

    tmp$p_NoSmoke <- p
    tmp$OR_NoSmoke <- OR

    tmp

},mc.cores=1))

result1 <- result1[order(result1$rate_njmu , decreasing = T),]
result1[is.na(result1)] <- 0


result1 <- result1[,
    c(
    "gene" ,
    c("freq_njmu" , "freq_tcga" , "rate_njmu" , "rate_tcga" , "OR_All" , "p_All") ,
    c("freq_nosmoke_njmu" , "freq_nosmoke_tcga" , "rate_nosmoke_njmu" , "rate_nosmoke_tcga" , "OR_NoSmoke" , "p_NoSmoke") ,
    c("freq_smoke_njmu" , "freq_smoke_tcga" , "rate_smoke_njmu" , "rate_smoke_tcga" , "OR_Smoke" , "p_Smoke")
    )]

out_name <- paste0(out_path , "/" , tumor_type , ".mutRate.compare.tsv")
write.table(result1,out_name,sep="\t" , quote =F , row.names = F )


###########################################################################################

plotRate <- function( result_use = result_use , images_name = images_name , title = title , width = width , height = height ){
    ###########################################################################################
    ## 柱状图展示
    ## p值
    result_use$p_text=ifelse(result_use$p>=0.05,"","*")
    result_use$p_text=ifelse(result_use$p<0.05 & result_use$p>0.01,"*",result_use$p_text)
    result_use$p_text=ifelse(result_use$p<0.01 & result_use$p>0.001,"**",result_use$p_text)
    result_use$p_text=ifelse(result_use$p<0.001 ,"***",result_use$p_text)
    gene_order <- result_use$gene

    ##########################################################################################

    result2 <- melt(result_use[,c("gene","njmu_rate","tcga_rate","p_text")])
    result2$value_percent=paste(round(result2$value * 100),"%",sep="")
    result2$gene <- factor( result2$gene , levels = gene_order , order = T )
    result2$variable <- ifelse( result2$variable == "njmu_rate" , "NJMU" , "TCGA" )

    p <- ggplot(result2,mapping = aes(variable,value,fill=variable))+geom_bar(stat='identity',position='stack') + facet_grid(.~gene)+
    theme_bw() +
    labs(title = title) +
    theme(strip.text.x = element_text(size = 11))+
    theme(panel.grid=element_blank())+labs(x = 'Genes',y = 'Mutation Rate') +
    theme(axis.title =element_text(size = 15),axis.text =element_text(size = 14, color = 'black'))+
    theme(axis.ticks.x = element_blank(),axis.text.x = element_blank())+
    geom_text(aes(label=value_percent, y=value+0.003), position=position_dodge(0.9), vjust=0)+
    geom_text(aes(label=p_text, y=0.7,x=1.5),size=7)+
    xlab("") +
    ylim(0,1) +
    theme(
      title =element_text(size=4, face='bold'),
      legend.title = element_blank(),
      legend.text = element_text(size = 12),
      legend.key.width = unit(1, "cm"),
      legend.key.height = unit(1, "cm"),
      plot.title = element_text(size = 30, face = "bold")
    )

    ggsave( images_name , p , width = width , height = height )

}

#### 所有的
result_use <- result1[,c("gene","rate_njmu","rate_tcga" , "p_All")]
colnames(result_use) <- c("gene","njmu_rate","tcga_rate" , "p")

title <- "All"
width <- 22
height <- 7
images_name <- paste0(out_path , "/MutRate.compareTCGA.All.pdf")
plotRate( result_use = result_use , images_name = images_name , title = title , width = width , height = height )

## 只显示差异的
result_use <- subset( result_use , p < 0.05)
title <- "All_Diff"
width <- nrow(result_use)
height <- 7
images_name <- paste0(out_path , "/MutRate.compareTCGA.All.Diff.pdf")
plotRate( result_use = result_use , images_name = images_name , title = title , width = width , height = height )

#### 吸烟人群中
result_use <- result1[,c("gene","rate_smoke_njmu","rate_smoke_tcga" , "p_Smoke")]
colnames(result_use) <- c("gene","njmu_rate","tcga_rate" , "p")
images_name <- paste0(out_path , "/MutRate.compareTCGA.Smoke.pdf")

title <- "Smoke"
width <- 22
height <- 7
plotRate( result_use = result_use , images_name = images_name , title = title , width = width , height = height )

## 只显示差异的
result_use <- subset( result_use , p < 0.05)
title <- "Smoke_Diff"
width <- 10
height <- 7
images_name <- paste0(out_path , "/MutRate.compareTCGA.Smoke.Diff.pdf")
plotRate( result_use = result_use , images_name = images_name , title = title , width = width , height = height )



#### 非吸烟人群中
result_use <- result1[,c("gene","rate_nosmoke_njmu","rate_nosmoke_tcga" , "p_NoSmoke")]
colnames(result_use) <- c("gene","njmu_rate","tcga_rate" , "p")
images_name <- paste0(out_path , "/MutRate.compareTCGA.NoSmoke.pdf")

title <- "NoSmoke"
width <- 20
height <- 7
plotRate( result_use = result_use , images_name = images_name , title = title , width = width , height = height )

## 只显示差异的
result_use <- subset( result_use , p < 0.05)
title <- "NoSmoke_Diff"
width <- 5
height <- 7
images_name <- paste0(out_path , "/MutRate.compareTCGA.NoSmoke.Diff.pdf")
plotRate( result_use = result_use , images_name = images_name , title = title , width = width , height = height )


###########################################################################################
## 20220420
## 展示所有，上面为非吸烟人群，下面为吸烟人群

result2 <- result1[,c("gene" , "rate_nosmoke_njmu" , "rate_nosmoke_tcga" , "OR_NoSmoke" , "p_NoSmoke")]
result3 <- result1[,c("gene" , "rate_smoke_njmu" , "rate_smoke_tcga" , "OR_Smoke" , "p_Smoke")]

result2$Type <- "Non-Smoker"
result3$Type <- "Smoker"

colnames(result2) <- c("gene" , "njmu_rate" , "tcga_rate" , "OR" , "p" , "Type")
colnames(result3) <- c("gene" , "njmu_rate" , "tcga_rate" , "OR" , "p" , "Type")

result4 <- rbind( result2 , result3 )

###########################################################################################
## 柱状图展示
## p值
result_use <- result4
result_use$p_text=ifelse(result_use$p>=0.05,"","*")
result_use$p_text=ifelse(result_use$p<0.05 & result_use$p>0.01,"*",result_use$p_text)
result_use$p_text=ifelse(result_use$p<0.01 & result_use$p>0.001,"**",result_use$p_text)
result_use$p_text=ifelse(result_use$p<0.001 ,"***",result_use$p_text)

gene_order <- unique(result_use$gene)

##########################################################################################

result2 <- melt(result_use[,c("gene","njmu_rate","tcga_rate","p_text","Type")])
result2$p_pos <- 0.65
result2$percent_pos <- result2$value + 0.003

result2$value_percent=paste(round(result2$value * 100),"%",sep="")
result2$gene <- factor( result2$gene , levels = gene_order , order = T )
result2$variable <- ifelse( result2$variable == "njmu_rate" , "NJMU" , "TCGA" )

result2$value[result2$Type == "Smoker"] <- -result2$value[result2$Type == "Smoker"]
result2$p_pos[result2$Type == "Smoker"] <- -result2$p_pos[result2$Type == "Smoker"]
result2$percent_pos[result2$Type == "Smoker"] <- -result2$percent_pos[result2$Type == "Smoker"] - 0.03


p <- ggplot(result2,mapping = aes(x = gene , y = value , fill = variable)) + geom_bar(stat = 'identity', position = 'dodge') + 
facet_grid(vars(Type) , scales = "free")+
theme_bw() +
theme(axis.text.y = element_blank())+
theme(panel.grid=element_blank())+labs(x = 'Genes',y = 'Mutation Rate') +
theme(axis.title =element_text(size = 15),axis.text =element_text(size = 12, color = 'black'))+
geom_text(aes(label= value_percent , x = gene, y= percent_pos ), position=position_dodge(0.9), vjust=0 ,color="black", face='bold' , family="Helvetica")+
geom_text(aes(label=p_text , x = gene , y = p_pos ),size=7,family="Helvetica")+
theme(axis.text.x = element_text(angle = 45, vjust = 1, hjust = 1 , size = 15 , family="Helvetica"))+
xlab("") +
theme(
  title =element_text(size=4, face='bold'),
  legend.title = element_blank(),
  legend.text = element_text(size = 12),
  legend.key.width = unit(1, "cm"),
  legend.key.height = unit(1, "cm"),
  plot.title = element_text(size = 30, face = "bold"),
  text = element_text(family="Helvetica")
)

width <- 18
height <- 7
images_name <- paste0(out_path , "/MutRate.compareTCGA.All.Gene.pdf")
ggsave( images_name , p , width = width , height = height )

## 图例展示基因
show_gene <- result2[abs(result2$value) > 0.2,]$gene
result3 <- subset( result2 , gene %in% show_gene )

col = c(rgb(red=230,green=74,blue=53,alpha=255,max=255),
        rgb(red=0,green=159,blue=134,alpha=255,max=255),
        rgb(red=77,green=186,blue=212,alpha=255,max=255),
        "grey"
)


p <- ggplot(data = result3 , mapping = aes(x = gene , y = value , fill = factor(variable)) ) + 
geom_bar(stat = 'identity', position = 'dodge') + 
scale_fill_manual( values = c(col[1] , col[3]) ) +
facet_grid(vars(Type) , scales = "free")+
theme_bw() +
theme(axis.text.y = element_blank())+
theme(panel.grid=element_blank())+labs(x = 'Genes',y = 'Mutation Rate') +
theme(axis.title =element_text(size = 15),axis.text =element_text(size = 24, color = 'black'))+
geom_text(aes(label= value_percent , x = gene, y= percent_pos ), position=position_dodge(0.9), size = 10, vjust=0 ,color="black", face='bold' , family="Helvetica")+
geom_text(aes(label=p_text , x = gene , y = p_pos ),size=7,family="Helvetica")+
theme(axis.text.x = element_text(angle = 45, vjust = 1, hjust = 1 , size = 30 , family="Helvetica"))+
xlab("") +
theme(
  title =element_text(size=21, face='bold'),
  legend.title = element_blank(),
  legend.text = element_text(size = 12),
  legend.key.width = unit(1, "cm"),
  legend.key.height = unit(1, "cm"),
  plot.title = element_text(size = 30, face = "bold"),
  text = element_text(family="Helvetica")
)



width <- 15
height <- 16
images_name <- paste0(out_path , "/MutRate.compareTCGA.All.Gene.grid.pdf")
ggsave( images_name , p , width = width , height = height )

